"""
Cicero district and rep lookup.
"""
# XXX: unicode()s due to "can't adapt" error in psycopg2, which
# doesn't like <class 'suds.sax.text.Text'>

import datetime
import re
import suds

# add parent directory to the path
import os.path as p
import sys
sys.path.insert(0, p.dirname(p.dirname(p.abspath(__file__))))

# import local settings
import settings
import django.core.management as m
m.setup_environ(settings)

import django.db.models as dbm
import ourcity.reps.models as models

def get_token():
    token_url = 'http://cicero.azavea.com/Azavea.Cicero.WebService.v2/AuthenticationService.asmx?WSDL'
    token_client = suds.client.Client(token_url)
    token = token_client.service.GetToken(userName=settings.CICERO_USERNAME, 
        password=settings.CICERO_PASSWORD)
    return token

class NoAddressMatch(Exception):
    def __init__(self):
        Exception.__init__(self, 
            'No match found for that address. Please try another.')


class CiceroConnection(object):
    def __init__(self):
        self.token = get_token()

    def get_cicero_location(self, address, city, state, country, 
            district_type):
        geocode_url = 'http://cicero.azavea.com/Azavea.Cicero.WebService.v2/GeocodingService.asmx?WSDL'
        geocode_client = suds.client.Client(geocode_url)
        locations = geocode_client.service.GetDistrictsByAddress(
            authToken=self.token,
            address=address, 
            city=city, 
            state=state,  # Two-letter abbrev. works
            postalCode=None,  # assume no postal code input
            country=country,  # US for United States
            districtType=district_type,  # e.g., LOCAL, STATE_UPPER, ALL
        )
        try:
            first_location = locations.Location[0]
        except AttributeError:
            first_location = None
        return first_location

    def query_address(self, address, city, state, country='US',
            district_type='ALL'):
        cicero_location = self.get_cicero_location(address, city, state,
                country, district_type)
        if not cicero_location:
            return cicero_location, []
        cicero_districts = cicero_location.Districts.DistrictInfo or []
        districts = []
        for cicero_district in cicero_districts:
            try:
                district = models.District.objects.get(
                    name=unicode(cicero_district.DistrictID), 
                    type__name=unicode(cicero_district.DistrictType), 
                    city=unicode(cicero_district.City or ''),
                    state=unicode(cicero_district.State or ''),
                    country=unicode(cicero_district.Country),
                )
            except models.District.DoesNotExist:
                continue  # ignore missing districts
            districts.append(district)
        return cicero_location, districts

    def create_districts_for_address(self, address, city, state, country='US',
            district_type='ALL'):
        cicero_location = self.get_cicero_location(address, city, state,
                country, district_type)
        if not cicero_location:
            raise NoAddressMatch
        cicero_districts = cicero_location.Districts.DistrictInfo or []
        districts = []
        for cicero_district in cicero_districts:
            type, created = models.DistrictType.objects.get_or_create(
                name=unicode(cicero_district.DistrictType))
            subtype, created = models.DistrictSubType.objects.get_or_create(
                name=unicode(cicero_district.DistrictSubType))
            district, created = models.District.objects.get_or_create(
                name=unicode(cicero_district.DistrictID), 
                type=type, 
                city=unicode(cicero_district.City or ''),
                state=unicode(cicero_district.State or ''),
                country=unicode(cicero_district.Country),
                defaults={'subtype': subtype},
            )
            district_str = district.type.name + district.city + \
                district.state + district.name
            print 'Found %s, Created: %s' % (district_str, created)
            
    def get_reps_for_address(self, address, city, state, country='US', 
            district_type='ALL'):
        "Direct call to Cicero for testing purposes."
        officials_query_url = 'http://cicero.azavea.com/Azavea.Cicero.WebService.v2/ElectedOfficialQueryService.asmx?WSDL'
        query_client = suds.client.Client(officials_query_url)
        cicero_reps = query_client.service.GetOfficialsByAddress(
            authToken=self.token,
            address=address,
            city=city,
            state=state,
            country=country,
            districtType=district_type,
            includeAtLarge=True,
        )
        return cicero_reps

    def get_reps_for_district(self, district_id):
        "For testing"
        district = models.District.objects.get(id=district_id)
        officials_query_url = 'http://cicero.azavea.com/Azavea.Cicero.WebService.v2/ElectedOfficialQueryService.asmx?WSDL'
        query_client = suds.client.Client(officials_query_url)
        cicero_reps = query_client.service.GetOfficialsByDistrictIDList(
            authToken=self.token,
            districtIDList=district.name, 
            city=district.city, 
            state=district.state,
            country=district.country,
            districtType=district.type,
            includeAtLarge=True,
        )
        #reps = [(x.FirstName, x.LastName) for x in cicero_reps.ElectedOfficialInfo]
        #return reps
        return cicero_reps

    def update_all_reps(self):
        districts = models.District.objects.all()
        district_queries = {}
        for district in districts:
            # group districts by type, city, and state
            agg = district.reps.aggregate(dbm.Max('last_update_date'))
            update_date = agg['last_update_date__max']
            query_id = district.type.name + district.city + district.state
            print '%s %s updated %s' % (query_id, district.name, update_date)
            try:
                query = district_queries[query_id]
            except KeyError:
                district_queries[query_id] = dict(
                    district_type=district.type.name,
                    city=district.city,
                    state=district.state,
                    update_dates=[update_date],
                    districts=[district],
                )
            else:
                query['update_dates'].append(update_date)
                query['districts'].append(district)
        officials_query_url = 'http://cicero.azavea.com/Azavea.Cicero.WebService.v2/ElectedOfficialQueryService.asmx?WSDL'
        query_client = suds.client.Client(officials_query_url)
        for query_id in district_queries:
            query = district_queries[query_id]
            cicero_query = dict(
                authToken=self.token, 
                city=query['city'],
                state=query['state'], 
                districtType=query['district_type'], 
                districtIDList=','.join(x.name for x in query['districts']),
                includeAtLarge=(query['district_type'] == 'LOCAL'),
            )
            update_date = query_client.service.GetLastDistrictUpdateDate(
                    **cicero_query)
            needs_update = False
            if None in query['update_dates']:
                print 'New district found.  Updating %s ...' % query_id
                needs_update = True
            elif update_date > max(query['update_dates']):
                print '%s is more recent than %s.  Updating %s ...' % \
                        (update_date, max(query['update_dates']), query_id)
                needs_update = True
            else:
                print '%s is not more recent than %s.  Not updating %s ...' % \
                        (update_date, max(query['update_dates']), query_id)
            if needs_update:
                # disassociate reps from affected districts
                for district in query['districts']:
                    district.reps.clear()
                c_reps = query_client.service.GetOfficialsByDistrictIDList(
                        **cicero_query)
                for cicero_rep in c_reps.ElectedOfficialInfo:
                    self.update_rep(cicero_rep)

    def update_rep(self, cicero_rep):
        try:
            print 'Updating %s %s ...' % (cicero_rep.FirstName, cicero_rep.LastName)
        except AttributeError:
            print 'Updating rep with missing name ...'
        try:
            cicero_party = unicode(cicero_rep.Party)
        except AttributeError:
            party = None
        else:
            try:
                party = models.Party.objects.get(name=cicero_party)
            except models.Party.DoesNotExist:
                party = models.Party(name=cicero_party, 
                    initial=cicero_party[0])
                party.save()
        try:
            cicero_title = unicode(cicero_rep.Title)
        except AttributeError:
            office = None
        else:
            cicero_title = cicero_title.split(',')[0]
            try:
                office = models.Office.objects.get(name=cicero_title)
            except models.Office.DoesNotExist:
                office = models.Office(name=cicero_title)
                office.save()
        try:
            city = unicode(cicero_rep.RepresentingCity)
        except AttributeError:
            city = ''
        try:
            state = unicode(cicero_rep.RepresentingState)
        except AttributeError:
            state = ''
        # district should already exist because we pull reps by district
        try:
            district = models.District.objects.get(
                name=unicode(cicero_rep.DistrictID), 
                type__name=unicode(cicero_rep.DistrictType), 
                city=city,
                state=state,
                country=unicode(cicero_rep.Country),
            )
        except models.District.DoesNotExist:
            print cicero_rep.DistrictID
            print cicero_rep.DistrictType
            print city
            print state
            print cicero_rep.Country
            raise
        rep_fields = dict(
            elected_official_id=unicode(cicero_rep.ElectedOfficialID),
            district_id=district.id,
            party_id=party.id if party else None,
            office_id=office.id if office else None,
        )
        skip_list = [
            'DistrictID',
            'DistrictIDLabel',
            'DistrictType',
            'ElectedOfficialID',
            'IsAtLarge',
            'Party',
            'RepresentingCity',
            'RepresentingState',
        ]
        cap_words = re.compile('[A-Z][a-z0-9]*')
        now = datetime.datetime.now()
        for key, value in cicero_rep:
            if key in skip_list:
                continue
            # don't split EMail1 and EMail2 into e_mail1 and e_mail2
            if key.startswith('EMail'):
                key = key.lower()
            # turn KeyNames into key_names
            field_name = '_'.join([x.lower() for x in cap_words.findall(key)])
            if field_name and value:
                if type(value) == type(now):
                    rep_fields[field_name] = value
                else:
                    rep_fields[field_name] = unicode(value)
        try:
            rep = models.Representative.objects.get(
                    elected_official_id=unicode(cicero_rep.ElectedOfficialID))
        except models.Representative.DoesNotExist:
            rep = models.Representative()
        rep.__dict__.update(rep_fields)
        rep.save()
        return rep

    def get_credits_remaining(self):
        account_query_url = 'http://cicero.azavea.com/Azavea.Cicero.WebService.v2/AccountInfoService.asmx?WSDL'
        query_client = suds.client.Client(account_query_url)
        credits = query_client.service.GetCreditsRemaining(
            authToken=self.token)
        return credits


if __name__ == '__main__':
    cc = CiceroConnection()
    if sys.argv[1] == 'query':
        location, districts = cc.query_address(sys.argv[2], 'Philadelphia', 'PA')
        print location
    elif sys.argv[1] == 'create-districts':
        print 'Creating districts for address ...'
        cc.create_districts_for_address(sys.argv[2], 'Philadelphia', 'PA')
    elif sys.argv[1] == 'get-reps-by-address':
        print 'Pulling reps from Cicero ...'
        print cc.get_reps_for_address(sys.argv[2], 'Philadelphia', 'PA')
    elif sys.argv[1] == 'get-reps-by-district':
        print 'Pulling reps from Cicero ...'
        #for rep in cc.get_reps_for_district(sys.argv[2]):
        #    print '%s %s' % rep
        print cc.get_reps_for_district(sys.argv[2])
    elif sys.argv[1] == 'update-reps':
        print 'Updating reps ...'
        cc.update_all_reps()
    elif sys.argv[1] == 'credits-remaining':
        print 'Remaining credits: %s' % cc.get_credits_remaining()
    else:
        print 'No valid option selected.'

